import os
import json
from collections import defaultdict

tm_lst = [
    '', # agents model path
]

ego_lst = [
    '', # agents model path
]

pair_lst = [[ego, tm] for tm in tm_lst for ego in ego_lst]
test_name = 'test'
for i, pair in enumerate(pair_lst):
    print(f'{i+1} / {len(pair_lst)}')
    os.system(f"bash src/scripts/test.sh {pair[0]} {pair[1]} {test_name}")

env = 'lbf'

def calculate_desired_food_ratio(data):
    food = defaultdict(int)
    n_get_food_episode = len([x for x in data['food_lst'] if x])
    for k, v in data['res'].items():
        if k == '':
            continue
        if 'A' in k:
            food['A'] += v
        if 'B' in k:
            food['B'] += v
        if 'C' in k:
            food['C'] += v
        if 'D' in k:
            food['D'] += v
        if 'E' in k:
            food['E'] += v
    if 'lbf4_A_' in data['tm']:
        n_good_food = food['A']
    elif 'lbf4_B_' in data['tm']:
        n_good_food = food['B']
    elif 'lbf4_C_' in data['tm']:
        n_good_food = food['C']
    elif 'lbf4_D_' in data['tm']:
        n_good_food = food['D']
    elif 'lbf4_E_' in data['tm']:
        n_good_food = food['E']
    else:
        assert 0
    if n_get_food_episode == 0:
        return 0
    else:
        return n_good_food / len(data['food_lst'])

def calculate_desired_prey_ratio(data):
    prey = defaultdict(int)
    n_caught_prey_episode = len([x for x in data['food_lst'] if x])
    for k, v in data['res'].items():
        if k == '':
            continue
        if '0' in k:
            prey['S1'] += v
        elif '1' in k:
            prey['S2'] += v
        elif k == '23':
            prey['R1R2'] = v
        elif k == '24':
            prey['R1R3'] = v
        elif k == '34':
            prey['R2R3'] = v
    if 'human1' in data['tm']:
        n_good_prey = prey['S1']
    elif 'human2' in data['tm']:
        n_good_prey = prey['S2']
    elif 'human3' in data['tm']:
        n_good_prey = prey['R1R2']
    elif 'human4' in data['tm']:
        n_good_prey = prey['R1R3']
    elif 'human5' in data['tm']:
        n_good_prey = prey['R2R3']
    else:
        assert 0
    if n_caught_prey_episode == 0:
        return 0
    else:
        return n_good_prey / len(data['food_lst'])

def calculate_kill_ratio(data):
    kill = 0
    for k, v in data['res'].items():
        if k != '':
            kill += v
    return kill / len(data['food_lst'])

def calculate_desired_kill_ratio(data):
    kill = defaultdict(int)
    for k, v in data['res'].items():
        if k == '':
            continue
        if 'A' in k:
            kill['A'] += v
        if 'B' in k:
            kill['B'] += v
        if 'C' in k:
            kill['C'] += v
        if 'D' in k:
            kill['D'] += v
        if 'E' in k:
            kill['E'] += v
    if '_A_' in data['tm']:
        n_good_kill = kill['A']
    elif '_B_' in data['tm']:
        n_good_kill = kill['B']
    elif '_C_' in data['tm']:
        n_good_kill = kill['C']
    elif '_D_' in data['tm']:
        n_good_kill = kill['D']
    elif '_E_' in data['tm']:
        n_good_kill = kill['E']
    else:
        print(data['tm'])
        assert 0
    return n_good_kill / len(data['food_lst'])

if env == 'lbf':
    all_res_path = ""
    calculate_desired_ratio = calculate_desired_food_ratio
if env == 'pp':
    all_res_path = ""
    calculate_desired_ratio = calculate_desired_prey_ratio
if env == 'sc2':
    all_res_path = ""
    calculate_desired_ratio = calculate_desired_kill_ratio
name2res = {}
res_lst = sorted(os.listdir(all_res_path))
for res in res_lst:
    if not os.path.isdir(os.path.join(all_res_path, res)):
        continue
    if 'res.json' in os.listdir(os.path.join(all_res_path, res)):
        res_path = os.path.join(all_res_path, res, 'res.json')
        with open(res_path) as f:
            data = json.load(f)
        desired_ratio = calculate_desired_ratio(data)
        if 'ego' not in data:
            seed = res.split('_')[1]
            data['ego'] = f'models/prompt_{seed}'
        name2res[f"{data['ego'].split('models/')[1]}+++{data['tm'].split('models/')[1]}"] = {
            'return_original': data['ret'] if env == 'lbf' or env == 'pp' else calculate_kill_ratio(data),
            'desired_ratio': desired_ratio
        }
    elif 'res_0.json' in os.listdir(os.path.join(all_res_path, res)):
        for res_file in os.listdir(os.path.join(all_res_path, res)):
            if not ('res_' in res_file and '.json' in res_file):
                continue
            head_id = res_file.split('_')[1].split('.')[0]
            res_path = os.path.join(all_res_path, res, res_file)
            with open(res_path) as f:
                data = json.load(f)
            desired_ratio = calculate_desired_ratio(data)
            name2res[f"{data['ego'].split('models/')[1]}---{head_id}+++{data['tm'].split('models/')[1]}"] = {
                'return_original': data['ret'] if env == 'lbf' or env == 'pp' else calculate_kill_ratio(data),
                'desired_ratio': desired_ratio
            }

with open(os.path.join(all_res_path, "res.json"), "w") as file:
    json.dump(name2res, file)
